import numpy as np
import scipy.io as sio
import os
import os.path as osp

data_dir = '../tid2013/'

file_root = data_dir + 'distorted_images' + '/'
list_file = data_dir + 'mos_with_names' + '.txt'


filename = [line.rstrip('\n') for line in open(
            osp.join(list_file))]
S_name = [] 
scores = []
for i in filename:
    S_name.append(i.split()[1])
    scores.append(float(i.split()[0]))
            
ref = S_name

TotalNum = len(ref)

# Create a shuffle of reference images (1-24)
num_ref_images = 24
shuff = np.random.permutation(range(1, num_ref_images + 1))

# Calculate split sizes based on number of reference images
Num_tr = int(num_ref_images * 0.7)  # 70% for training
Num_val = int(num_ref_images * 0.2)  # 20% for validation
Num_te = num_ref_images - Num_tr - Num_val  # 10% for testing

# Ensure output directory exists
os.makedirs('../score_tid2013', exist_ok=True)

# Create output files
train_file = open('../score_tid2013/ft_tid2013'+'_train.txt', "w")
val_file = open('../score_tid2013/ft_tid2013'+'_val.txt', "w")
test_file = open('../score_tid2013/ft_tid2013'+'_test.txt', "w")

# Shuffle all images
shuff_txt = np.random.permutation(range(TotalNum))    

# Write training set
for i in shuff_txt:
    # Check if the reference image number is in the training set
    ref_img_num = int(ref[i][1:3])
    if ref_img_num in shuff[:Num_tr]:
        rel_path = data_dir + 'distorted_images/' + ref[i]
        abs_path = os.path.abspath(rel_path)
        labels = scores[i]
        train_file.write('%s %6.2f\n' % (abs_path, labels))

# Write validation set
for i in shuff_txt:
    ref_img_num = int(ref[i][1:3])
    if ref_img_num in shuff[Num_tr:Num_tr+Num_val]:
        rel_path = data_dir + 'distorted_images/' + ref[i]
        abs_path = os.path.abspath(rel_path)
        labels = scores[i]
        val_file.write('%s %6.2f\n' % (abs_path, labels))

# Write test set
for i in shuff_txt:
    ref_img_num = int(ref[i][1:3])
    if ref_img_num in shuff[Num_tr+Num_val:]:
        rel_path = data_dir + 'distorted_images/' + ref[i]
        abs_path = os.path.abspath(rel_path)
        labels = scores[i]
        test_file.write('%s %6.2f\n' % (abs_path, labels))

train_file.close()
val_file.close()
test_file.close()
print("处理完成！数据已按照7:2:1的比例分为训练集、验证集和测试集")